package com.walle.ak47.commons.jaxrs.interceptor.log;

import java.io.FilterWriter;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.io.Writer;

import org.apache.cxf.interceptor.Fault;
import org.apache.cxf.interceptor.LoggingMessage;
import org.apache.cxf.interceptor.LoggingOutInterceptor;
import org.apache.cxf.interceptor.StaxOutInterceptor;
import org.apache.cxf.io.CacheAndWriteOutputStream;
import org.apache.cxf.io.CachedOutputStream;
import org.apache.cxf.io.CachedOutputStreamCallback;
import org.apache.cxf.message.Message;
import org.apache.cxf.phase.Phase;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;

public class Log4jLoggingOutJsonInterceptor extends AbstractLog4jLoggingInterceptor {

	protected static Logger logger = Logger.getLogger(Log4jLoggingOutJsonInterceptor.class.getPackage().getName());

	private static final String LOG_SETUP = LoggingOutInterceptor.class.getName() + ".log-setup";

	public Log4jLoggingOutJsonInterceptor(String phase) {
		super(phase);
		addBefore(StaxOutInterceptor.class.getName());
	}

	public Log4jLoggingOutJsonInterceptor() {
		this(Phase.PRE_STREAM);
	}

	public Log4jLoggingOutJsonInterceptor(int lim) {
		this();
		limit = lim;
	}

	public Log4jLoggingOutJsonInterceptor(PrintWriter w) {
		this();
		this.writer = w;
	}

	public void handleMessage(Message message) throws Fault {
		final OutputStream os = message.getContent(OutputStream.class);
		final Writer iowriter = message.getContent(Writer.class);
		if (os == null && iowriter == null) {
			return;
		}
		// Logger cLogger = getMessageLogger(message);
		Logger cLogger = this.getLogger();
		if (cLogger.getLevel().equals(Level.INFO) || writer != null) {
			// Write the output while caching it for the log message
			boolean hasLogged = message.containsKey(LOG_SETUP);
			if (!hasLogged) {
				message.put(LOG_SETUP, Boolean.TRUE);
				if (os != null) {
					final CacheAndWriteOutputStream newOut = new CacheAndWriteOutputStream(os);
					if (threshold > 0) {
						newOut.setThreshold(threshold);
					}
					if (limit > 0) {
						newOut.setCacheLimit(limit);
					}
					message.setContent(OutputStream.class, newOut);
					newOut.registerCallback(new LoggingCallback(cLogger, message, os));
				} else {
					message.setContent(Writer.class, new LogWriter(cLogger, message, iowriter));
				}
			}
		}
	}

	private JsonLoggingMessage setupBuffer(Message message) {
		String id = (String) message.getExchange().get(LoggingMessage.ID_KEY);
		if (id == null) {
			id = LoggingMessage.nextId();
			message.getExchange().put(LoggingMessage.ID_KEY, id);
		}
		final JsonLoggingMessage buffer = new JsonLoggingMessage("ResponseMessage", id);

		Integer responseCode = (Integer) message.get(Message.RESPONSE_CODE);
		if (responseCode != null) {
			buffer.getResponseCode().append(responseCode);
		}

		String encoding = (String) message.get(Message.ENCODING);
		if (encoding != null) {
			buffer.getEncoding().append(encoding);
		}
		String httpMethod = (String) message.get(Message.HTTP_REQUEST_METHOD);
		if (httpMethod != null) {
			buffer.getHttpMethod().append(httpMethod);
		}
		String address = (String) message.get(Message.ENDPOINT_ADDRESS);
		if (address != null) {
			buffer.getAddress().append(address);
			String uri = (String) message.get(Message.REQUEST_URI);
			if (uri != null && !address.startsWith(uri)) {
				if (!address.endsWith("/") && !uri.startsWith("/")) {
					buffer.getAddress().append("/");
				}
				buffer.getAddress().append(uri);
			}
		}
		String ct = (String) message.get(Message.CONTENT_TYPE);
		if (ct != null) {
			buffer.getContentType().append(ct);
		}
		Object headers = message.get(Message.PROTOCOL_HEADERS);
		if (headers != null) {
			buffer.getHeader().append(headers);
		}
		return buffer;
	}

	private class LogWriter extends FilterWriter {
		StringWriter out2;
		int count;
		Logger logger; // NOPMD
		Message message;
		final int lim;

		public LogWriter(Logger logger, Message message, Writer writer) {
			super(writer);
			this.logger = logger;
			this.message = message;
			if (!(writer instanceof StringWriter)) {
				out2 = new StringWriter();
			}
			lim = limit == -1 ? Integer.MAX_VALUE : limit;
		}

		public void write(int c) throws IOException {
			super.write(c);
			if (out2 != null && count < lim) {
				out2.write(c);
			}
			count++;
		}

		public void write(char[] cbuf, int off, int len) throws IOException {
			super.write(cbuf, off, len);
			if (out2 != null && count < lim) {
				out2.write(cbuf, off, len);
			}
			count += len;
		}

		public void write(String str, int off, int len) throws IOException {
			super.write(str, off, len);
			if (out2 != null && count < lim) {
				out2.write(str, off, len);
			}
			count += len;
		}

		public void close() throws IOException {
			JsonLoggingMessage buffer = setupBuffer(message);
			if (count >= lim) {
				buffer.getMessage().append("(message truncated to " + lim + " bytes)\n");
			}
			StringWriter w2 = out2;
			if (w2 == null) {
				w2 = (StringWriter) out;
			}
			String ct = (String) message.get(Message.CONTENT_TYPE);
			try {
				writePayload(buffer.getPayload(), w2, ct);
			} catch (Exception ex) {
				// ignore
			}
			log(logger, buffer.toString());
			message.setContent(Writer.class, out);
			super.close();
		}
	}

	protected String formatLoggingMessage(JsonLoggingMessage buffer) {
		return buffer.toString();
	}

	class LoggingCallback implements CachedOutputStreamCallback {

		private final Message message;
		private final OutputStream origStream;
		private final Logger logger; // NOPMD
		private final int lim;

		public LoggingCallback(final Logger logger, final Message msg, final OutputStream os) {
			this.logger = logger;
			this.message = msg;
			this.origStream = os;
			this.lim = limit == -1 ? Integer.MAX_VALUE : limit;
		}

		public void onFlush(CachedOutputStream cos) {

		}

		public void onClose(CachedOutputStream cos) {
			JsonLoggingMessage buffer = setupBuffer(message);

			String ct = (String) message.get(Message.CONTENT_TYPE);
			if (!isShowBinaryContent() && isBinaryContent(ct)) {
				buffer.getMessage().append(BINARY_CONTENT_MESSAGE).append('\n');
				log(logger, formatLoggingMessage(buffer));
				return;
			}

			if (cos.getTempFile() == null) {
				// buffer.append("Outbound Message:\n");
				if (cos.size() >= lim) {
					buffer.getMessage().append("(message truncated to " + lim + " bytes)\n");
				}
			} else {
				buffer.getMessage().append("Outbound Message (saved to tmp file):\n");
				buffer.getMessage().append("Filename: " + cos.getTempFile().getAbsolutePath() + "\n");
				if (cos.size() >= lim) {
					buffer.getMessage().append("(message truncated to " + lim + " bytes)\n");
				}
			}
			try {
				String encoding = (String) message.get(Message.ENCODING);
				writePayload(buffer.getPayload(), cos, encoding, ct);
			} catch (Exception ex) {
				// ignore
			}

			log(logger, formatLoggingMessage(buffer));
			try {
				// empty out the cache
				cos.lockOutputStream();
				cos.resetOut(null, false);
			} catch (Exception ex) {
				// ignore
			}
			message.setContent(OutputStream.class, origStream);
		}
	}

	protected Logger getLogger() {
		if (null == logger.getLevel()) {
			logger.setLevel(Level.INFO);
		}
		return logger;
	}

}